"""
主窗口GUI模块

该模块实现了应用程序的主窗口界面。
主要功能：
1. 视频文件加载和显示
2. 目标选择和跟踪控制
3. 参数设置和调整
4. 结果保存和导出

作者：目标跟踪系统
日期：2024
"""

import sys
import os
import cv2
import numpy as np
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from typing import Optional, Tuple
import time

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from src.core.tracker import FusionTracker
from src.utils.video_processor import VideoProcessor
from src.utils.visualization import TrackingVisualizer
from .video_widget import VideoWidget
from .control_panel import ControlPanel


class MainWindow(QMainWindow):
    """
    主窗口类
    
    应用程序的主界面，集成视频显示、控制面板和跟踪功能。
    提供完整的用户交互界面。
    """
    
    def __init__(self):
        super().__init__()
        
        # 初始化组件
        self.tracker = FusionTracker()
        self.video_processor = VideoProcessor()
        self.visualizer = TrackingVisualizer()
        
        # 界面状态
        self.is_tracking = False
        self.is_playing = False
        self.current_frame = None
        self.target_bbox = None
        
        # 定时器
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        
        # 初始化界面
        self.init_ui()
        
    def init_ui(self):
        """
        初始化用户界面
        """
        self.setWindowTitle("MeanShift-Kalman目标跟踪系统")
        self.setGeometry(100, 100, 1200, 800)
        
        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        main_layout = QHBoxLayout(central_widget)
        
        # 创建视频显示组件
        self.video_widget = VideoWidget()
        self.video_widget.target_selected.connect(self.on_target_selected)
        
        # 创建控制面板
        self.control_panel = ControlPanel()
        self.control_panel.load_video_clicked.connect(self.load_video)
        self.control_panel.start_tracking_clicked.connect(self.start_tracking)
        self.control_panel.stop_tracking_clicked.connect(self.stop_tracking)
        self.control_panel.play_pause_clicked.connect(self.toggle_play_pause)
        self.control_panel.parameters_changed.connect(self.update_parameters)
        
        # 布局设置
        main_layout.addWidget(self.video_widget, 3)
        main_layout.addWidget(self.control_panel, 1)
        
        # 创建菜单栏
        self.create_menu_bar()
        
        # 创建状态栏
        self.create_status_bar()
        
    def create_menu_bar(self):
        """
        创建菜单栏
        """
        menubar = self.menuBar()
        
        # 文件菜单
        file_menu = menubar.addMenu('文件')
        
        load_action = QAction('加载视频', self)
        load_action.setShortcut('Ctrl+O')
        load_action.triggered.connect(self.load_video)
        file_menu.addAction(load_action)
        
        save_action = QAction('保存结果', self)
        save_action.setShortcut('Ctrl+S')
        save_action.triggered.connect(self.save_result)
        file_menu.addAction(save_action)
        
        file_menu.addSeparator()
        
        exit_action = QAction('退出', self)
        exit_action.setShortcut('Ctrl+Q')
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)
        
        # 跟踪菜单
        track_menu = menubar.addMenu('跟踪')
        
        start_action = QAction('开始跟踪', self)
        start_action.setShortcut('Space')
        start_action.triggered.connect(self.start_tracking)
        track_menu.addAction(start_action)
        
        stop_action = QAction('停止跟踪', self)
        stop_action.setShortcut('Esc')
        stop_action.triggered.connect(self.stop_tracking)
        track_menu.addAction(stop_action)
        
        # 帮助菜单
        help_menu = menubar.addMenu('帮助')
        
        about_action = QAction('关于', self)
        about_action.triggered.connect(self.show_about)
        help_menu.addAction(about_action)
        
    def create_status_bar(self):
        """
        创建状态栏
        """
        self.status_bar = self.statusBar()
        
        # 状态标签
        self.status_label = QLabel("就绪")
        self.status_bar.addWidget(self.status_label)
        
        # FPS标签
        self.fps_label = QLabel("FPS: 0")
        self.status_bar.addPermanentWidget(self.fps_label)
        
        # 帧数标签
        self.frame_label = QLabel("帧: 0/0")
        self.status_bar.addPermanentWidget(self.frame_label)
        
    def load_video(self):
        """
        加载视频文件
        """
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择视频文件", "", 
            "视频文件 (*.mp4 *.avi *.mov *.mkv *.wmv);;所有文件 (*)")
        
        if file_path:
            if self.video_processor.open_video(file_path):
                self.status_label.setText(f"已加载视频: {os.path.basename(file_path)}")
                
                # 显示第一帧
                ret, frame = self.video_processor.read_frame()
                if ret:
                    self.current_frame = frame
                    self.video_widget.update_frame(frame)
                    self.update_frame_info()
                    
                # 更新控制面板
                video_info = self.video_processor.get_video_info()
                self.control_panel.update_video_info(video_info)
                
                # 重置跟踪状态
                self.reset_tracking()
            else:
                QMessageBox.warning(self, "错误", "无法加载视频文件")
                
    def on_target_selected(self, bbox: Tuple[int, int, int, int]):
        """
        目标选择回调
        
        参数：
            bbox: 选择的目标边界框
        """
        self.target_bbox = bbox
        self.status_label.setText("目标已选择，可以开始跟踪")
        
        # 初始化跟踪器
        if self.current_frame is not None:
            self.tracker.initialize(self.current_frame, bbox)
            self.control_panel.enable_tracking_controls(True)
            
    def start_tracking(self):
        """
        开始跟踪
        """
        if self.target_bbox is None:
            QMessageBox.warning(self, "警告", "请先选择目标区域")
            return
            
        if not self.video_processor.cap:
            QMessageBox.warning(self, "警告", "请先加载视频文件")
            return
            
        self.is_tracking = True
        self.is_playing = True
        self.status_label.setText("跟踪中...")
        
        # 启动定时器
        self.timer.start(33)  # 约30fps
        
        # 更新控制面板
        self.control_panel.set_tracking_state(True)
        
    def stop_tracking(self):
        """
        停止跟踪
        """
        self.is_tracking = False
        self.is_playing = False
        self.timer.stop()
        
        self.status_label.setText("跟踪已停止")
        self.control_panel.set_tracking_state(False)
        
    def toggle_play_pause(self):
        """
        切换播放/暂停状态
        """
        if self.is_playing:
            self.is_playing = False
            self.timer.stop()
            self.status_label.setText("已暂停")
        else:
            self.is_playing = True
            self.timer.start(33)
            self.status_label.setText("播放中...")
            
    def update_frame(self):
        """
        更新帧显示
        """
        if not self.video_processor.cap:
            return
            
        start_time = time.time()
        
        # 读取下一帧
        ret, frame = self.video_processor.read_frame()
        if not ret:
            self.stop_tracking()
            return
            
        self.current_frame = frame
        
        # 执行跟踪
        if self.is_tracking:
            success, bbox = self.tracker.track(frame)
            
            # 获取跟踪信息
            tracking_info = self.tracker.get_tracking_info()
            
            # 计算FPS
            frame_time = time.time() - start_time
            fps = self.visualizer.calculate_fps(frame_time)
            tracking_info['fps'] = fps
            
            # 创建可视化结果
            result_frame = self.visualizer.create_result_image(
                frame, bbox, tracking_info, show_trajectory=True)
            
            # 更新显示
            self.video_widget.update_frame(result_frame)
            
            # 更新控制面板
            self.control_panel.update_tracking_info(tracking_info)
            
            # 更新状态栏
            self.fps_label.setText(f"FPS: {fps:.1f}")
            
        else:
            self.video_widget.update_frame(frame)
            
        # 更新帧信息
        self.update_frame_info()
        
    def update_frame_info(self):
        """
        更新帧信息显示
        """
        if self.video_processor.cap:
            current = self.video_processor.current_frame_index
            total = self.video_processor.frame_count
            self.frame_label.setText(f"帧: {current}/{total}")
            
    def update_parameters(self, params: dict):
        """
        更新跟踪参数
        
        参数：
            params: 参数字典
        """
        self.tracker.set_parameters(
            similarity_threshold=params.get('similarity_threshold'),
            max_lost_frames=params.get('max_lost_frames'),
            kalman_weight=params.get('kalman_weight')
        )
        
    def save_result(self):
        """
        保存跟踪结果
        """
        if self.current_frame is None:
            QMessageBox.warning(self, "警告", "没有可保存的结果")
            return
            
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存结果", "", "图像文件 (*.png *.jpg *.bmp)")
        
        if file_path:
            if self.visualizer.save_result(self.current_frame, file_path):
                QMessageBox.information(self, "成功", "结果已保存")
            else:
                QMessageBox.warning(self, "错误", "保存失败")
                
    def reset_tracking(self):
        """
        重置跟踪状态
        """
        self.is_tracking = False
        self.is_playing = False
        self.target_bbox = None
        self.timer.stop()
        
        self.tracker.reset()
        self.visualizer.reset_trajectory()
        self.control_panel.reset()
        
        self.status_label.setText("就绪")
        
    def show_about(self):
        """
        显示关于对话框
        """
        QMessageBox.about(self, "关于", 
                         "MeanShift-Kalman目标跟踪系统\n\n"
                         "基于MeanShift算法和Kalman滤波器的融合跟踪系统\n"
                         "版本: 1.0\n"
                         "作者: 目标跟踪系统开发团队")
        
    def closeEvent(self, event):
        """
        窗口关闭事件
        """
        self.video_processor.close_video()
        event.accept() 